from Data.DataStorage import HDFS
from Util.Process import INFO
from Util.SQL import select
from Util.Return import *
import json


# 完成模型存储相关的任务: HDFS上传、下载；模型上星
class ModelFS:
    def __init__(self, model_id):
        self.model_id = model_id
        # self.state_list = ['upload', 'validate', 'compress', 'evaluate', 'ready']
        self.hdfs = HDFS()

    def pull_model(self, state, sub_path=''):
        if state not in ['compress', 'evaluate', 'ready']:
            INFO("ERROR", f"ModelTransfer pull_model Error: '{state}' is not true state.")
        else:
            json_string = select(f"select format from model where model_id = '{self.model_id}'")[0]['format']
            format_json = json.loads(json_string.replace("'", '"'))
            if state == 'compress':
                local_compress_path = get_path(os.path.abspath(__file__), 2, get_config("model", "compress_path"))
                if sub_path != '':
                    local_compress_path = os.path.join(local_compress_path, sub_path)
                hdfs_raw_path = "/" + get_config("model", "raw_path") + "/" + self.model_id + "." + format_json["raw"]
                os.makedirs(local_compress_path, exist_ok=True)
                self.hdfs.pull(hdfs_raw_path, local_compress_path)
            elif state == 'evaluate':
                local_evaluate_path = get_path(os.path.abspath(__file__), 2, get_config("model", "evaluate_path"))
                hdfs_ready_path = "/" + get_config("model", "ready_path") + "/" + self.model_id + "." + format_json["ready"]
                if sub_path != '':
                    local_evaluate_path = os.path.join(local_evaluate_path, sub_path)
                os.makedirs(local_evaluate_path, exist_ok=True)
                self.hdfs.pull(hdfs_ready_path, local_evaluate_path)

    def push_model(self, state, sub_path=''):
        if state not in ['validate', 'compress']:
            INFO("ERROR", f"ModelTransfer push_model Error: '{state}' is not true state.")
        else:
            json_string = select(f"select format from model where model_id = '{self.model_id}'")[0]['format']
            format_json = json.loads(json_string.replace("'", '"'))
            if state == 'validate':
                local_raw_path = os.path.join(get_path(os.path.abspath(__file__), 2, get_config("model", "raw_path")), self.model_id + "." + format_json["raw"])
                hdfs_raw_path = "/" + get_config("model", "raw_path")
                if self.hdfs.exist(hdfs_raw_path):
                    self.hdfs.delete(hdfs_raw_path)
                self.hdfs.push(hdfs_raw_path, local_raw_path)
            elif state == 'compress':
                local_compress_path = get_path(os.path.abspath(__file__), 2, get_config("model", "compress_path"))
                if sub_path != '':
                    local_compress_path = os.path.join(local_compress_path, sub_path)
                print(sub_path)
                local_compress_path = os.path.join(local_compress_path, self.model_id + "." + format_json["ready"])
                hdfs_ready_path = "/" + get_config("model", "ready_path")
                if self.hdfs.exist(hdfs_ready_path):
                    self.hdfs.delete(hdfs_ready_path)
                self.hdfs.push(hdfs_ready_path, local_compress_path)
